"""
Mem0记忆管理接口
基于Mem0构建的记忆管理系统，支持AI智能体的长期记忆存储和检索
"""

import asyncio
import os
from datetime import datetime
from typing import Dict, Any, List, Optional

from dotenv import load_dotenv
from mem0 import Memory, MemoryClient
from sqlalchemy.orm import Session

from config.logging_config import get_logger
from pojo.ai_web_memory_his_model import AIWebMemoryHisModel

# 加载环境变量
load_dotenv()


class Mem0Interface:
    """Mem0记忆管理接口类"""

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        初始化Mem0接口
        
        Args:
            config: Mem0配置，包括向量数据库、LLM等配置
        """
        self.mem0_client = None
        self.logger = get_logger(__name__)
        self._initialization_error = None
        self.config = self._get_default_config()

        try:
            # OSS
            # self.mem0_client = Memory.from_config(self.config)
            # 远程
            self.mem0_client = MemoryClient()
            self.logger.info("✅ Mem0客户端初始化成功")

        except Exception as e:
            self._initialization_error = f"Mem0客户端初始化失败: {e}"
            self.logger.error(f"❌ {self._initialization_error}")
            self.logger.error(f"配置信息: {self.config}")

    def _get_default_config(self) -> Dict[str, Any]:
        """获取默认Mem0配置"""
        # 确保数据库目录存在
        db_path = os.path.abspath("./catch/mem0")
        try:
            os.makedirs(db_path, exist_ok=True)
            # 检查目录是否可写
            test_file = os.path.join(db_path, ".test_write")
            with open(test_file, 'w') as f:
                f.write("test")
            os.remove(test_file)
        except Exception as e:
            self.logger.warning(f"数据库目录创建或写入测试失败: {e}")
            # 回退到系统临时目录
            import tempfile
            db_path = os.path.join(tempfile.gettempdir(), "mem0_cache")
            os.makedirs(db_path, exist_ok=True)
        
        return {
            "vector_store": {
                "provider": "chroma",
                "config": {
                    "collection_name": "ai_web_memories",
                    "path": db_path,
                }
            },
            "llm": {
                "provider": "deepseek",
                "config": {
                    "api_key": os.getenv("DEEPSEEK_API_KEY", ""),
                    "model": "deepseek-chat",
                }
            },
            "history_db_path": os.path.join(db_path, "history.db"),
            "version": "v1.1",
        }

    async def add_memory(
            self,
            content: str,
            user_id: str,
            agent_id: str,
            scene_id: str,
            session_id: str,
            metadata: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        添加记忆到Mem0
        
        Args:
            content: 记忆内容
            user_id: 用户ID
            agent_id: 智能体ID
            scene_id: 场景ID
            session_id: 会话ID
            metadata: 额外元数据
            
        Returns:
            记忆添加结果
        """
        try:
            if not self.mem0_client:
                return {"success": False, "error": "Mem0客户端未初始化"}

            # 准备记忆数据
            memory_data = {
                "messages": [{"role": "user", "content": content}],
                "user_id": user_id,
                "agent_id": agent_id,
                "metadata": {
                    "scene_id": scene_id,
                    "session_id": session_id,
                    "timestamp": datetime.now().isoformat(),
                    **(metadata or {})
                }
            }

            # 添加到Mem0
            result = await asyncio.get_event_loop().run_in_executor(
                None,
                lambda: self.mem0_client.add(**memory_data)
            )

            self.logger.info(f"记忆添加成功: user_id={user_id}, agent_id={agent_id}")
            return {"success": True, "result": result}

        except Exception as e:
            self.logger.error(f"添加记忆失败: {e}")
            return {"success": False, "error": str(e)}

    def search_memories(
            self,
            query: str,
            user_id: str,
            agent_id: Optional[str] = None,
            scene_id: Optional[str] = None,
            limit: int = 10
    ) -> List[Dict[str, Any]]:
        """
        搜索相关记忆（同步方法）
        
        Args:
            query: 搜索查询
            user_id: 用户ID
            agent_id: 智能体ID（可选）
            scene_id: 场景ID（可选）
            limit: 返回结果数量限制
            
        Returns:
            相关记忆列表
        """
        try:
            if not self.mem0_client:
                self.logger.warning("Mem0客户端未初始化")
                return []

            # 构建搜索参数
            search_params = {
                "query": query,
                "user_id": user_id,
                "limit": limit
            }

            if agent_id:
                search_params["agent_id"] = agent_id

            self.logger.info(f"开始搜索记忆: query={query}, user_id={user_id}, params={search_params}")

            # 执行同步搜索
            results = self.mem0_client.search(**search_params)
            
            self.logger.info(f"Mem0原始搜索结果数量: {len(results) if results else 0}")

            # 过滤结果（如果需要按scene_id筛选）
            if scene_id and results:
                filtered_results = []
                for result in results:
                    metadata = result.get("metadata", {})
                    if metadata.get("scene_id") == scene_id:
                        filtered_results.append(result)
                results = filtered_results
                self.logger.info(f"按scene_id={scene_id}过滤后结果数量: {len(results)}")

            self.logger.info(f"记忆搜索完成: 找到{len(results)}条相关记忆")
            return results

        except Exception as e:
            self.logger.error(f"搜索记忆失败: {e}", exc_info=True)
            return []

    async def get_user_memories(
            self,
            user_id: str,
            agent_id: Optional[str] = None,
            scene_id: Optional[str] = None,
            limit: int = 50
    ) -> List[Dict[str, Any]]:
        """
        获取用户的所有记忆
        
        Args:
            user_id: 用户ID
            agent_id: 智能体ID（可选）
            scene_id: 场景ID（可选）
            limit: 返回结果数量限制
            
        Returns:
            用户记忆列表
        """
        try:
            if not self.mem0_client:
                return []

            # 构建获取参数
            get_params = {
                "user_id": user_id,
                "limit": limit
            }

            if agent_id:
                get_params["agent_id"] = agent_id

            # 获取记忆
            results = await asyncio.get_event_loop().run_in_executor(
                None,
                lambda: self.mem0_client.get_all(**get_params)
            )

            # 按scene_id过滤（如果指定）
            if scene_id and results:
                filtered_results = []
                for result in results:
                    metadata = result.get("metadata", {})
                    if metadata.get("scene_id") == scene_id:
                        filtered_results.append(result)
                results = filtered_results

            return results

        except Exception as e:
            self.logger.error(f"获取用户记忆失败: {e}")
            return []

    async def update_memory(
            self,
            memory_id: str,
            content: str,
            metadata: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        更新记忆内容
        
        Args:
            memory_id: 记忆ID
            content: 新的记忆内容
            metadata: 新的元数据
            
        Returns:
            更新结果
        """
        try:
            if not self.mem0_client:
                return {"success": False, "error": "Mem0客户端未初始化"}

            # 准备更新数据
            update_data = {
                "memory_id": memory_id,
                "data": content
            }

            if metadata:
                update_data["metadata"] = metadata

            # 执行更新
            result = await asyncio.get_event_loop().run_in_executor(
                None,
                lambda: self.mem0_client.update(**update_data)
            )

            return {"success": True, "result": result}

        except Exception as e:
            self.logger.error(f"更新记忆失败: {e}")
            return {"success": False, "error": str(e)}

    async def delete_memory(self, memory_id: str) -> bool:
        """
        删除记忆
        
        Args:
            memory_id: 记忆ID
            
        Returns:
            删除是否成功
        """
        try:
            if not self.mem0_client:
                return False

            # 执行删除
            await asyncio.get_event_loop().run_in_executor(
                None,
                lambda: self.mem0_client.delete(memory_id)
            )

            return True

        except Exception as e:
            self.logger.error(f"删除记忆失败: {e}")
            return False

    async def delete_memories_by_metadata(
            self,
            user_id: str,
            agent_id: Optional[str] = None,
            scene_id: Optional[str] = None,
            session_id: Optional[str] = None
    ) -> bool:
        """
        根据metadata删除记忆
        
        Args:
            user_id: 用户ID
            agent_id: 智能体ID（可选）
            scene_id: 场景ID（可选）
            session_id: 会话ID（可选）
            
        Returns:
            删除是否成功
        """
        try:
            if not self.mem0_client:
                return False

            # 构建查询条件获取要删除的记忆
            memories = await self.get_user_memories(user_id, agent_id, scene_id)

            deleted_count = 0
            for memory in memories:
                metadata = memory.get("metadata", {})

                # 检查是否匹配删除条件
                should_delete = True

                if session_id and metadata.get("session_id") != session_id:
                    should_delete = False

                if should_delete:
                    memory_id = memory.get("id")
                    if memory_id:
                        try:
                            if await self.delete_memory(memory_id):
                                deleted_count += 1
                        except Exception as e:
                            self.logger.warning(f"删除记忆 {memory_id} 失败: {e}")

            self.logger.info(f"按metadata删除记忆完成: 删除了{deleted_count}条记忆")
            return deleted_count > 0

        except Exception as e:
            self.logger.error(f"按metadata删除记忆失败: {e}")
            return False

    async def get_memory_ids_by_session(
            self,
            user_id: str,
            agent_id: str,
            scene_id: str,
            session_id: str
    ) -> List[str]:
        """
        获取指定会话的所有记忆ID
        
        Args:
            user_id: 用户ID
            agent_id: 智能体ID
            scene_id: 场景ID
            session_id: 会话ID
            
        Returns:
            记忆ID列表
        """
        try:
            memories = await self.get_user_memories(user_id, agent_id, scene_id)

            memory_ids = []
            for memory in memories:
                metadata = memory.get("metadata", {})
                if metadata.get("session_id") == session_id:
                    memory_id = memory.get("id")
                    if memory_id:
                        memory_ids.append(memory_id)

            return memory_ids

        except Exception as e:
            self.logger.error(f"获取会话记忆ID失败: {e}")
            return []

    def is_available(self) -> bool:
        """检查Mem0是否可用"""
        is_available = self.mem0_client is not None

        if not is_available and self._initialization_error:
            self.logger.warning(f"⚠️ Mem0不可用: {self._initialization_error}")

        return is_available

    def get_initialization_error(self) -> Optional[str]:
        """获取初始化错误信息"""
        return self._initialization_error

    async def get_memory_stats(
            self,
            user_id: str,
            agent_id: Optional[str] = None,
            scene_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        获取记忆统计信息
        
        Args:
            user_id: 用户ID
            agent_id: 智能体ID（可选）
            scene_id: 场景ID（可选）
            
        Returns:
            记忆统计信息
        """
        try:
            memories = await self.get_user_memories(user_id, agent_id, scene_id)

            stats = {
                "total_memories": len(memories),
                "agent_distribution": {},
                "scene_distribution": {},
                "recent_memories": 0
            }

            # 统计各智能体和场景的记忆分布
            recent_threshold = datetime.now().timestamp() - 7 * 24 * 3600  # 7天前

            for memory in memories:
                metadata = memory.get("metadata", {})

                # 智能体分布
                agent = metadata.get("agent_id", "unknown")
                stats["agent_distribution"][agent] = stats["agent_distribution"].get(agent, 0) + 1

                # 场景分布
                scene = metadata.get("scene_id", "unknown")
                stats["scene_distribution"][scene] = stats["scene_distribution"].get(scene, 0) + 1

                # 最近记忆统计
                timestamp = metadata.get("timestamp")
                if timestamp:
                    try:
                        memory_time = datetime.fromisoformat(timestamp.replace('Z', '+00:00'))
                        if memory_time.timestamp() > recent_threshold:
                            stats["recent_memories"] += 1
                    except:
                        pass

            return stats

        except Exception as e:
            self.logger.error(f"获取记忆统计失败: {e}")
            return {"total_memories": 0, "error": str(e)}


class Mem0MemoryProvider:
    """Mem0记忆提供者，为不同组件提供统一的记忆接口"""

    def __init__(self, mem0_interface: Mem0Interface, db_session: Session):
        """
        初始化记忆提供者
        
        Args:
            mem0_interface: Mem0接口实例
            db_session: 数据库会话
        """
        self.mem0 = mem0_interface
        self.db = db_session
        self.logger = get_logger(__name__)

    async def store_conversation_memory(
            self,
            user_message: str,
            assistant_message: str,
            user_id: str,
            agent_id: str,
            scene_id: str,
            session_id: str,
            metadata: Optional[Dict[str, Any]] = None
    ) -> bool:
        """
        存储对话记忆（同时存储到Mem0和数据库）
        
        Args:
            user_message: 用户消息
            assistant_message: 助手回复
            user_id: 用户ID
            agent_id: 智能体ID
            scene_id: 场景ID
            session_id: 会话ID
            metadata: 额外元数据
            
        Returns:
            存储是否成功
        """
        try:
            # 构建对话内容
            conversation_content = f"用户: {user_message}\n助手: {assistant_message}"

            # 存储到Mem0
            if self.mem0.is_available():
                mem0_result = await self.mem0.add_memory(
                    content=conversation_content,
                    user_id=user_id,
                    agent_id=agent_id,
                    scene_id=scene_id,
                    session_id=session_id,
                    metadata=metadata
                )

                if not mem0_result.get("success"):
                    self.logger.warning(f"Mem0存储失败: {mem0_result.get('error')}")

            # 存储到数据库
            memory_record = AIWebMemoryHisModel(
                scene_id=scene_id,
                agent_id=agent_id,
                session_id=session_id,
                user_id=user_id,
                his_text=conversation_content,
                status=1,  # 临时状态
                importance_level=3,  # 默认重要程度
                create_time=datetime.now()
            )

            if metadata:
                memory_record.set_extra_info(metadata)

            self.db.add(memory_record)
            self.db.commit()

            self.logger.info(f"对话记忆存储成功: session_id={session_id}")
            return True

        except Exception as e:
            self.logger.error(f"存储对话记忆失败: {e}")
            self.db.rollback()
            return False

    async def retrieve_context_memories(
            self,
            query: str,
            user_id: str,
            agent_id: str,
            scene_id: str,
            limit: int = 5
    ) -> List[str]:
        """
        检索上下文相关的记忆
        
        Args:
            query: 查询内容
            user_id: 用户ID
            agent_id: 智能体ID
            scene_id: 场景ID
            limit: 返回数量限制
            
        Returns:
            相关记忆内容列表
        """
        try:
            memories = []

            # 从Mem0检索
            if self.mem0.is_available():
                mem0_memories = await self.mem0.search_memories(
                    query=query,
                    user_id=user_id,
                    agent_id=agent_id,
                    scene_id=scene_id,
                    limit=limit
                )

                for memory in mem0_memories:
                    content = memory.get("memory", memory.get("content", ""))
                    if content:
                        memories.append(content)

            # 如果Mem0结果不足，从数据库补充
            if len(memories) < limit:
                remaining_limit = limit - len(memories)

                db_memories = self.db.query(AIWebMemoryHisModel).filter(
                    AIWebMemoryHisModel.scene_id == scene_id,
                    AIWebMemoryHisModel.agent_id == agent_id,
                    AIWebMemoryHisModel.user_id == user_id,
                    AIWebMemoryHisModel.his_text.like(f"%{query}%")
                ).order_by(AIWebMemoryHisModel.create_time.desc()).limit(remaining_limit).all()

                for db_memory in db_memories:
                    if db_memory.his_text not in memories:
                        memories.append(db_memory.his_text)

            return memories[:limit]

        except Exception as e:
            self.logger.error(f"检索上下文记忆失败: {e}")
            return []

    async def get_session_history(
            self,
            session_id: str,
            limit: int = 10
    ) -> List[Dict[str, Any]]:
        """
        获取会话历史记录
        
        Args:
            session_id: 会话ID
            limit: 返回数量限制
            
        Returns:
            会话历史记录列表
        """
        try:
            db_memories = self.db.query(AIWebMemoryHisModel).filter(
                AIWebMemoryHisModel.session_id == session_id
            ).order_by(AIWebMemoryHisModel.create_time.desc()).limit(limit).all()

            history = []
            for memory in reversed(db_memories):  # 按时间正序返回
                history.append({
                    "id": memory.id,
                    "content": memory.his_text,
                    "timestamp": memory.create_time.isoformat(),
                    "importance_level": memory.importance_level,
                    "status": memory.status
                })

            return history

        except Exception as e:
            self.logger.error(f"获取会话历史失败: {e}")
            return []

    async def delete_memories_by_metadata(
            self,
            user_id: str,
            agent_id: str,
            scene_id: str,
            session_id: str
    ) -> bool:
        """
        根据metadata删除记忆（Mem0和数据库）
        
        Args:
            user_id: 用户ID
            agent_id: 智能体ID
            scene_id: 场景ID
            session_id: 会话ID
            
        Returns:
            删除是否成功
        """
        try:
            success = True

            # 删除Mem0记忆
            if self.mem0.is_available():
                try:
                    mem0_success = await self.mem0.delete_memories_by_metadata(
                        user_id=user_id,
                        agent_id=agent_id,
                        scene_id=scene_id,
                        session_id=session_id
                    )
                    if not mem0_success:
                        self.logger.warning("Mem0记忆删除失败")
                        success = False
                except Exception as e:
                    self.logger.warning(f"Mem0记忆删除失败: {e}")
                    success = False

            # 删除数据库记忆
            try:
                deleted_count = self.db.query(AIWebMemoryHisModel).filter(
                    AIWebMemoryHisModel.user_id == user_id,
                    AIWebMemoryHisModel.agent_id == agent_id,
                    AIWebMemoryHisModel.scene_id == scene_id,
                    AIWebMemoryHisModel.session_id == session_id
                ).delete()

                self.db.commit()
                self.logger.info(f"数据库记忆删除成功: 删除了{deleted_count}条记录")

            except Exception as e:
                self.logger.error(f"数据库记忆删除失败: {e}")
                self.db.rollback()
                success = False

            return success

        except Exception as e:
            self.logger.error(f"按metadata删除记忆失败: {e}")
            return False

    async def get_memory_ids_by_session(
            self,
            user_id: str,
            agent_id: str,
            scene_id: str,
            session_id: str
    ) -> List[str]:
        """
        获取指定会话的所有记忆ID
        
        Args:
            user_id: 用户ID
            agent_id: 智能体ID
            scene_id: 场景ID
            session_id: 会话ID
            
        Returns:
            记忆ID列表
        """
        try:
            return await self.mem0.get_memory_ids_by_session(
                user_id=user_id,
                agent_id=agent_id,
                scene_id=scene_id,
                session_id=session_id
            )

        except Exception as e:
            self.logger.error(f"获取会话记忆ID失败: {e}")
            return []

    async def delete_memory_by_id(self, memory_id: str) -> bool:
        """
        根据ID删除记忆
        
        Args:
            memory_id: 记忆ID
            
        Returns:
            删除是否成功
        """
        try:
            if self.mem0.is_available():
                return await self.mem0.delete_memory(memory_id)
            return False

        except Exception as e:
            self.logger.error(f"根据ID删除记忆失败: {e}")
            return False
